package com.viknix.threadpool.manager.client.util;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

/**
 * @Author: Dongqi
 * @Date: 2021/11/24 9:49
 * @Version 1.0
 * @Description:
 */
@Slf4j
public class Limiter {

    private static final int CAPACITY_SIZE = 1000;
    private static final int LIMIT_TIME = 1000;
    private static Cache<String, RateLimiter> cache = CacheBuilder.newBuilder()
            .initialCapacity(CAPACITY_SIZE).expireAfterAccess(1, TimeUnit.MINUTES)
            .build();

    /**
     * qps 5
     */
    private static double limit = 5;

    static {
        try {
            String limitTimeStr = System
                    .getProperty("limitTime", String.valueOf(limit));
            limit = Double.parseDouble(limitTimeStr);
            log.info("limitTime:{}", limit);
        } catch (Exception e) {
            log.error("init limitTime fail", e);
        }
    }

    public static boolean isLimit(String accessKeyID) {
        RateLimiter rateLimiter = null;
        try {
            rateLimiter = cache.get(accessKeyID, new Callable<RateLimiter>() {
                @Override
                public RateLimiter call() throws Exception {
                    return RateLimiter.create(limit);
                }
            });
        } catch (ExecutionException e) {
            log.error("create limit fail", e);
        }
        if (rateLimiter != null && !rateLimiter.tryAcquire(LIMIT_TIME, TimeUnit.MILLISECONDS)) {
            log.error("access_key_id:{} limited", accessKeyID);
            return true;
        }
        return false;
    }
}
